#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
@file    : adaptive_light_algorithm.py
@brief   : 自适应光源算法模块 - 香橙派3B视觉系统
@author  : IPCS 2025电赛培训项目
@version : 1.0
@date    : 2025-07-15

@note    : 该模块实现自适应光源算法，根据环境光强自动调节图像参数，
           支持曝光控制、增益调节、白平衡，确保视觉稳定性。
           与STM32F407VET6主控通过串口通信。
"""

import cv2
import numpy as np
import time
import threading
import serial
import struct
import json
from typing import Dict, Tuple, Optional, List
from enum import Enum
from dataclasses import dataclass

# ========================= 私有类型定义 =========================

class CameraStatus(Enum):
    """摄像头状态枚举"""
    OK = 0
    INIT_ERROR = 1
    CAPTURE_ERROR = 2
    COMM_ERROR = 3

class LightMode(Enum):
    """光线模式枚举"""
    AUTO = 0      # 自动模式
    MANUAL = 1    # 手动模式
    PRESET = 2    # 预设模式

@dataclass
class CameraParams:
    """摄像头参数结构体"""
    exposure: int = -1        # 曝光值 (-1表示自动)
    gain: float = 1.0         # 增益值 (1.0-16.0)
    brightness: int = 50      # 亮度 (0-100)
    contrast: int = 50        # 对比度 (0-100)
    saturation: int = 50      # 饱和度 (0-100)
    white_balance: int = -1   # 白平衡 (-1表示自动)

@dataclass
class LightMetrics:
    """光线检测指标"""
    brightness: float = 0.0   # 亮度值 (0-255)
    contrast: float = 0.0     # 对比度值
    histogram_peak: int = 0   # 直方图峰值位置
    overexposure_ratio: float = 0.0  # 过曝比例
    underexposure_ratio: float = 0.0 # 欠曝比例

# ========================= 用户配置参数区 =========================

# 摄像头配置
CAMERA_INDEX = 0                    # 摄像头索引
CAMERA_WIDTH = 1920                 # 图像宽度
CAMERA_HEIGHT = 1080                # 图像高度
CAMERA_FPS = 30                     # 帧率

# 串口通信配置
UART_PORT = "/dev/ttyUSB0"          # 串口设备路径
UART_BAUDRATE = 115200              # 波特率
FRAME_HEADER = 0xAA55               # 帧头
FRAME_END = 0x55AA                  # 帧尾
MAX_DATA_LENGTH = 256               # 最大数据长度

# 自适应算法参数
TARGET_BRIGHTNESS = 120             # 目标亮度值 (0-255)
BRIGHTNESS_TOLERANCE = 15           # 亮度容忍度
ADAPTATION_SPEED = 0.3              # 自适应速度 (0.1-1.0)
MIN_EXPOSURE = 1                    # 最小曝光值
MAX_EXPOSURE = 10000                # 最大曝光值
MIN_GAIN = 1.0                      # 最小增益
MAX_GAIN = 16.0                     # 最大增益

# 光线检测阈值
OVEREXPOSURE_THRESHOLD = 250        # 过曝阈值
UNDEREXPOSURE_THRESHOLD = 15        # 欠曝阈值
OVEREXPOSURE_RATIO_LIMIT = 0.05     # 过曝比例限制 (5%)
UNDEREXPOSURE_RATIO_LIMIT = 0.05    # 欠曝比例限制 (5%)

# 白平衡参数
AUTO_WB_ENABLED = True              # 是否启用自动白平衡
WB_ADAPTATION_RATE = 0.1            # 白平衡适应速率

# 调试模式
DEBUG_MODE = False                  # 调试模式开关
SHOW_HISTOGRAM = False              # 显示直方图

class AdaptiveLightAlgorithm:
    """自适应光源算法类"""
    
    def __init__(self):
        """
        @brief 初始化自适应光源算法模块
        """
        self.camera = None
        self.serial_port = None
        self.is_running = False
        self.current_params = CameraParams()
        self.light_metrics = LightMetrics()
        self.light_mode = LightMode.AUTO
        
        # 线程相关
        self.capture_thread = None
        self.adaptation_thread = None
        self.comm_thread = None
        self.frame_lock = threading.Lock()
        self.current_frame = None
        
        # 统计相关
        self.frame_count = 0
        self.last_adaptation_time = 0
        
    def adaptive_light_init(self) -> CameraStatus:
        """
        @brief 初始化自适应光源算法模块
        @retval CameraStatus.OK: 初始化成功
        @retval CameraStatus.INIT_ERROR: 初始化失败
        @note  初始化摄像头、串口通信和算法参数
        """
        try:
            # 初始化摄像头
            self.camera = cv2.VideoCapture(CAMERA_INDEX)
            if not self.camera.isOpened():
                if DEBUG_MODE:
                    print("错误: 无法打开摄像头")
                return CameraStatus.INIT_ERROR
            
            # 设置摄像头参数
            self.camera.set(cv2.CAP_PROP_FRAME_WIDTH, CAMERA_WIDTH)
            self.camera.set(cv2.CAP_PROP_FRAME_HEIGHT, CAMERA_HEIGHT)
            self.camera.set(cv2.CAP_PROP_FPS, CAMERA_FPS)
            
            # 初始化串口
            try:
                self.serial_port = serial.Serial(
                    port=UART_PORT,
                    baudrate=UART_BAUDRATE,
                    timeout=1
                )
                if DEBUG_MODE:
                    print(f"串口已连接: {UART_PORT}")
            except Exception as e:
                if DEBUG_MODE:
                    print(f"串口连接失败: {e}")
                return CameraStatus.INIT_ERROR
            
            # 初始化摄像头参数
            self._init_camera_params()
            
            if DEBUG_MODE:
                print("自适应光源算法模块初始化成功")
            
            return CameraStatus.OK
            
        except Exception as e:
            if DEBUG_MODE:
                print(f"初始化失败: {e}")
            return CameraStatus.INIT_ERROR
    
    def adaptive_light_start(self) -> bool:
        """
        @brief 启动自适应光源算法
        @retval True: 启动成功
        @retval False: 启动失败
        @note  启动图像采集、自适应算法和通信线程
        """
        if self.is_running:
            return True
        
        try:
            self.is_running = True
            
            # 启动图像采集线程
            self.capture_thread = threading.Thread(target=self._capture_loop)
            self.capture_thread.daemon = True
            self.capture_thread.start()
            
            # 启动自适应算法线程
            self.adaptation_thread = threading.Thread(target=self._adaptation_loop)
            self.adaptation_thread.daemon = True
            self.adaptation_thread.start()
            
            # 启动通信线程
            self.comm_thread = threading.Thread(target=self._communication_loop)
            self.comm_thread.daemon = True
            self.comm_thread.start()
            
            if DEBUG_MODE:
                print("自适应光源算法已启动")
            
            return True
            
        except Exception as e:
            if DEBUG_MODE:
                print(f"启动失败: {e}")
            self.is_running = False
            return False
    
    def adaptive_light_stop(self) -> None:
        """
        @brief 停止自适应光源算法
        @note  停止所有线程并释放资源
        """
        self.is_running = False
        
        # 等待线程结束
        if self.capture_thread and self.capture_thread.is_alive():
            self.capture_thread.join(timeout=1)
        
        if self.adaptation_thread and self.adaptation_thread.is_alive():
            self.adaptation_thread.join(timeout=1)
        
        if self.comm_thread and self.comm_thread.is_alive():
            self.comm_thread.join(timeout=1)
        
        # 释放资源
        if self.camera:
            self.camera.release()
        
        if self.serial_port:
            self.serial_port.close()
        
        if DEBUG_MODE:
            print("自适应光源算法已停止")
    
    def adaptive_light_get_params(self) -> CameraParams:
        """
        @brief 获取当前摄像头参数
        @retval CameraParams: 当前参数结构体
        """
        return self.current_params
    
    def adaptive_light_set_params(self, params: CameraParams) -> bool:
        """
        @brief 设置摄像头参数
        @param params: 参数结构体
        @retval True: 设置成功
        @retval False: 设置失败
        @note  手动设置参数，将切换到手动模式
        """
        try:
            self.current_params = params
            self.light_mode = LightMode.MANUAL
            self._apply_camera_params()
            
            if DEBUG_MODE:
                print(f"参数已更新: 曝光={params.exposure}, 增益={params.gain}")
            
            return True
        except Exception as e:
            if DEBUG_MODE:
                print(f"参数设置失败: {e}")
            return False
    
    def adaptive_light_set_mode(self, mode: LightMode) -> bool:
        """
        @brief 设置光线模式
        @param mode: 光线模式
        @retval True: 设置成功
        @retval False: 设置失败
        """
        try:
            self.light_mode = mode
            if DEBUG_MODE:
                print(f"光线模式已切换为: {mode.name}")
            return True
        except Exception as e:
            if DEBUG_MODE:
                print(f"模式设置失败: {e}")
            return False
    
    def adaptive_light_get_metrics(self) -> LightMetrics:
        """
        @brief 获取光线检测指标
        @retval LightMetrics: 光线指标结构体
        """
        return self.light_metrics
    
    def adaptive_light_capture_frame(self) -> Optional[np.ndarray]:
        """
        @brief 获取当前帧图像
        @retval np.ndarray: 图像数据，None表示获取失败
        """
        with self.frame_lock:
            return self.current_frame.copy() if self.current_frame is not None else None
    
    def adaptive_light_send_data(self, cmd: int, data: bytes) -> bool:
        """
        @brief 向STM32发送数据
        @param cmd: 命令字
        @param data: 数据内容
        @retval True: 发送成功
        @retval False: 发送失败
        """
        if not self.serial_port or not self.serial_port.is_open:
            return False
        
        try:
            # 构建数据包
            packet = self._build_packet(cmd, data)
            self.serial_port.write(packet)
            return True
        except Exception as e:
            if DEBUG_MODE:
                print(f"数据发送失败: {e}")
            return False
    
    # ========================= 私有方法 =========================
    
    def _init_camera_params(self) -> None:
        """初始化摄像头参数"""
        if self.camera:
            # 设置自动曝光
            self.camera.set(cv2.CAP_PROP_AUTO_EXPOSURE, 0.25)
            
            # 设置初始参数
            self.current_params.brightness = 50
            self.current_params.contrast = 50
            self.current_params.saturation = 50
            self.current_params.gain = 1.0
            
            self._apply_camera_params()
    
    def _apply_camera_params(self) -> None:
        """应用摄像头参数"""
        if not self.camera:
            return
        
        try:
            # 设置曝光
            if self.current_params.exposure > 0:
                self.camera.set(cv2.CAP_PROP_AUTO_EXPOSURE, 0.25)
                self.camera.set(cv2.CAP_PROP_EXPOSURE, self.current_params.exposure)
            else:
                self.camera.set(cv2.CAP_PROP_AUTO_EXPOSURE, 0.75)
            
            # 设置增益
            self.camera.set(cv2.CAP_PROP_GAIN, self.current_params.gain)
            
            # 设置亮度、对比度、饱和度
            self.camera.set(cv2.CAP_PROP_BRIGHTNESS, self.current_params.brightness / 100.0)
            self.camera.set(cv2.CAP_PROP_CONTRAST, self.current_params.contrast / 100.0)
            self.camera.set(cv2.CAP_PROP_SATURATION, self.current_params.saturation / 100.0)
            
            # 设置白平衡
            if AUTO_WB_ENABLED and self.current_params.white_balance <= 0:
                self.camera.set(cv2.CAP_PROP_AUTO_WB, 1)
            else:
                self.camera.set(cv2.CAP_PROP_AUTO_WB, 0)
                if self.current_params.white_balance > 0:
                    self.camera.set(cv2.CAP_PROP_WB_TEMPERATURE, self.current_params.white_balance)
            
        except Exception as e:
            if DEBUG_MODE:
                print(f"参数应用失败: {e}")
    
    def _capture_loop(self) -> None:
        """图像采集循环"""
        while self.is_running:
            try:
                ret, frame = self.camera.read()
                if ret:
                    with self.frame_lock:
                        self.current_frame = frame
                    self.frame_count += 1
                else:
                    if DEBUG_MODE:
                        print("图像采集失败")
                    time.sleep(0.1)
            except Exception as e:
                if DEBUG_MODE:
                    print(f"采集循环错误: {e}")
                time.sleep(0.1)
    
    def _adaptation_loop(self) -> None:
        """自适应算法循环"""
        while self.is_running:
            try:
                if self.light_mode == LightMode.AUTO:
                    current_time = time.time()
                    
                    # 限制自适应频率 (最多每100ms执行一次)
                    if current_time - self.last_adaptation_time > 0.1:
                        self._perform_adaptation()
                        self.last_adaptation_time = current_time
                
                time.sleep(0.05)  # 20Hz更新频率
                
            except Exception as e:
                if DEBUG_MODE:
                    print(f"自适应循环错误: {e}")
                time.sleep(0.1)
    
    def _perform_adaptation(self) -> None:
        """执行自适应调节"""
        frame = self.adaptive_light_capture_frame()
        if frame is None:
            return
        
        # 计算光线指标
        self._calculate_light_metrics(frame)
        
        # 执行自适应调节
        if self._needs_adjustment():
            self._adjust_camera_params()
            self._apply_camera_params()
    
    def _calculate_light_metrics(self, frame: np.ndarray) -> None:
        """计算光线检测指标"""
        # 转换为灰度图像
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        
        # 计算亮度 (平均值)
        self.light_metrics.brightness = np.mean(gray)
        
        # 计算对比度 (标准差)
        self.light_metrics.contrast = np.std(gray)
        
        # 计算直方图
        hist = cv2.calcHist([gray], [0], None, [256], [0, 256])
        self.light_metrics.histogram_peak = np.argmax(hist)
        
        # 计算过曝和欠曝比例
        total_pixels = gray.shape[0] * gray.shape[1]
        overexposed_pixels = np.sum(gray >= OVEREXPOSURE_THRESHOLD)
        underexposed_pixels = np.sum(gray <= UNDEREXPOSURE_THRESHOLD)
        
        self.light_metrics.overexposure_ratio = overexposed_pixels / total_pixels
        self.light_metrics.underexposure_ratio = underexposed_pixels / total_pixels
        
        if DEBUG_MODE and SHOW_HISTOGRAM:
            self._display_histogram(hist)
    
    def _needs_adjustment(self) -> bool:
        """判断是否需要调节"""
        brightness_diff = abs(self.light_metrics.brightness - TARGET_BRIGHTNESS)
        
        return (brightness_diff > BRIGHTNESS_TOLERANCE or
                self.light_metrics.overexposure_ratio > OVEREXPOSURE_RATIO_LIMIT or
                self.light_metrics.underexposure_ratio > UNDEREXPOSURE_RATIO_LIMIT)
    
    def _adjust_camera_params(self) -> None:
        """调节摄像头参数"""
        # 计算亮度差异
        brightness_diff = self.light_metrics.brightness - TARGET_BRIGHTNESS
        
        # 调节曝光
        if abs(brightness_diff) > BRIGHTNESS_TOLERANCE:
            if self.current_params.exposure <= 0:
                self.current_params.exposure = 100  # 切换到手动模式
            
            # 计算调节量
            adjustment = -brightness_diff * ADAPTATION_SPEED
            new_exposure = self.current_params.exposure * (1 + adjustment / 100)
            
            # 限制范围
            self.current_params.exposure = max(MIN_EXPOSURE, 
                                             min(MAX_EXPOSURE, int(new_exposure)))
        
        # 处理过曝
        if self.light_metrics.overexposure_ratio > OVEREXPOSURE_RATIO_LIMIT:
            # 降低曝光或增益
            if self.current_params.exposure > MIN_EXPOSURE:
                self.current_params.exposure = max(MIN_EXPOSURE, 
                                                 int(self.current_params.exposure * 0.9))
            elif self.current_params.gain > MIN_GAIN:
                self.current_params.gain = max(MIN_GAIN, 
                                             self.current_params.gain * 0.9)
        
        # 处理欠曝
        if self.light_metrics.underexposure_ratio > UNDEREXPOSURE_RATIO_LIMIT:
            # 增加曝光或增益
            if self.current_params.exposure < MAX_EXPOSURE:
                self.current_params.exposure = min(MAX_EXPOSURE, 
                                                 int(self.current_params.exposure * 1.1))
            elif self.current_params.gain < MAX_GAIN:
                self.current_params.gain = min(MAX_GAIN, 
                                             self.current_params.gain * 1.1)
        
        if DEBUG_MODE:
            print(f"参数调节 - 亮度: {self.light_metrics.brightness:.1f}, "
                  f"曝光: {self.current_params.exposure}, "
                  f"增益: {self.current_params.gain:.2f}")
    
    def _communication_loop(self) -> None:
        """通信循环"""
        while self.is_running:
            try:
                # 定期发送状态信息给STM32
                if self.frame_count % 30 == 0:  # 每30帧发送一次
                    self._send_status_to_stm32()
                
                # 接收STM32的命令
                self._receive_commands_from_stm32()
                
                time.sleep(0.01)  # 100Hz通信频率
                
            except Exception as e:
                if DEBUG_MODE:
                    print(f"通信循环错误: {e}")
                time.sleep(0.1)
    
    def _send_status_to_stm32(self) -> None:
        """向STM32发送状态信息"""
        try:
            # 构建状态数据包
            status_data = {
                'brightness': self.light_metrics.brightness,
                'contrast': self.light_metrics.contrast,
                'exposure': self.current_params.exposure,
                'gain': self.current_params.gain,
                'mode': self.light_mode.value
            }
            
            # 转换为JSON字符串
            json_str = json.dumps(status_data)
            data = json_str.encode('utf-8')
            
            # 发送数据包
            self.adaptive_light_send_data(0x01, data)
            
        except Exception as e:
            if DEBUG_MODE:
                print(f"状态发送失败: {e}")
    
    def _receive_commands_from_stm32(self) -> None:
        """接收STM32的命令"""
        if not self.serial_port or not self.serial_port.is_open:
            return
        
        try:
            if self.serial_port.in_waiting > 0:
                data = self.serial_port.read(self.serial_port.in_waiting)
                self._parse_received_data(data)
        except Exception as e:
            if DEBUG_MODE:
                print(f"命令接收失败: {e}")
    
    def _parse_received_data(self, data: bytes) -> None:
        """解析接收到的数据"""
        # 简单的命令解析示例
        if len(data) >= 3:
            cmd = data[0]
            if cmd == 0x10:  # 设置模式命令
                mode_value = data[1]
                if mode_value < len(LightMode):
                    self.light_mode = LightMode(mode_value)
            elif cmd == 0x11:  # 设置参数命令
                # 解析参数数据
                pass
    
    def _build_packet(self, cmd: int, data: bytes) -> bytes:
        """构建数据包"""
        # 数据包格式: [帧头2字节][长度2字节][命令1字节][数据][校验1字节][帧尾2字节]
        packet = bytearray()
        
        # 帧头
        packet.extend(struct.pack('<H', FRAME_HEADER))
        
        # 长度
        data_length = len(data) + 1  # 包含命令字
        packet.extend(struct.pack('<H', data_length))
        
        # 命令
        packet.append(cmd)
        
        # 数据
        packet.extend(data)
        
        # 简单校验和
        checksum = sum(packet) & 0xFF
        packet.append(checksum)
        
        # 帧尾
        packet.extend(struct.pack('<H', FRAME_END))
        
        return bytes(packet)
    
    def _display_histogram(self, hist: np.ndarray) -> None:
        """显示直方图 (调试用)"""
        import matplotlib.pyplot as plt
        
        plt.figure(figsize=(10, 4))
        plt.plot(hist)
        plt.title('Image Histogram')
        plt.xlabel('Pixel Value')
        plt.ylabel('Count')
        plt.axvline(x=TARGET_BRIGHTNESS, color='r', linestyle='--', label='Target')
        plt.axvline(x=self.light_metrics.brightness, color='g', linestyle='-', label='Current')
        plt.legend()
        plt.pause(0.01)

# ========================= 用户实现接口 =========================

def user_camera_init_callback() -> bool:
    """
    @brief 摄像头初始化回调函数（用户可选实现）
    @retval True: 允许初始化
    @retval False: 阻止初始化
    @note   用户可在此函数中进行自定义初始化操作
    """
    # 用户自定义初始化代码
    return True

def user_frame_process_callback(frame: np.ndarray) -> np.ndarray:
    """
    @brief 图像处理回调函数（用户可选实现）
    @param frame: 输入图像
    @retval np.ndarray: 处理后的图像
    @note   用户可在此函数中进行自定义图像处理
    """
    # 用户自定义图像处理代码
    # 例如: 添加滤波、去噪等操作
    return frame

def user_adaptation_callback(metrics: LightMetrics, params: CameraParams) -> CameraParams:
    """
    @brief 自适应算法回调函数（用户可选实现）
    @param metrics: 当前光线指标
    @param params: 当前摄像头参数
    @retval CameraParams: 调整后的参数
    @note   用户可在此函数中实现自定义的自适应算法
    """
    # 用户自定义自适应算法
    return params

# ========================= 全局实例 =========================

# 全局算法实例
_adaptive_light_instance = None

def get_adaptive_light_instance() -> AdaptiveLightAlgorithm:
    """
    @brief 获取自适应光源算法实例
    @retval AdaptiveLightAlgorithm: 算法实例
    @note   单例模式，确保全局只有一个实例
    """
    global _adaptive_light_instance
    if _adaptive_light_instance is None:
        _adaptive_light_instance = AdaptiveLightAlgorithm()
    return _adaptive_light_instance

# ========================= 简化接口函数 =========================

def init() -> bool:
    """简化的初始化接口"""
    instance = get_adaptive_light_instance()
    return instance.adaptive_light_init() == CameraStatus.OK

def start() -> bool:
    """简化的启动接口"""
    instance = get_adaptive_light_instance()
    return instance.adaptive_light_start()

def stop() -> None:
    """简化的停止接口"""
    instance = get_adaptive_light_instance()
    instance.adaptive_light_stop()

def get_frame() -> Optional[np.ndarray]:
    """简化的获取帧接口"""
    instance = get_adaptive_light_instance()
    return instance.adaptive_light_capture_frame()

def get_params() -> CameraParams:
    """简化的获取参数接口"""
    instance = get_adaptive_light_instance()
    return instance.adaptive_light_get_params()

def set_params(params: CameraParams) -> bool:
    """简化的设置参数接口"""
    instance = get_adaptive_light_instance()
    return instance.adaptive_light_set_params(params)

def get_metrics() -> LightMetrics:
    """简化的获取指标接口"""
    instance = get_adaptive_light_instance()
    return instance.adaptive_light_get_metrics()

if __name__ == "__main__":
    """测试代码"""
    print("自适应光源算法模块测试")
    
    # 初始化
    if not init():
        print("初始化失败")
        exit(1)
    
    # 启动算法
    if not start():
        print("启动失败")
        exit(1)
    
    try:
        # 运行测试
        for i in range(100):
            frame = get_frame()
            if frame is not None:
                metrics = get_metrics()
                params = get_params()
                
                print(f"帧 {i}: 亮度={metrics.brightness:.1f}, "
                      f"曝光={params.exposure}, 增益={params.gain:.2f}")
                
                # 显示图像 (可选)
                if DEBUG_MODE:
                    cv2.imshow('Camera', frame)
                    if cv2.waitKey(1) & 0xFF == ord('q'):
                        break
            
            time.sleep(0.1)
    
    except KeyboardInterrupt:
        print("用户中断")
    
    finally:
        # 停止算法
        stop()
        cv2.destroyAllWindows()
        print("测试完成")